#!/usr/bin/env python
import os
import sys
import subprocess
import base64
# import json
import argparse
from PIL import Image

compression_configs= {
    "ffmpeg": {
        "default": {
            "crf": 28,
            "preset": "slow",
            "audio_bitrate": "64k",
            "scale": "40:40",
            "quality": 80
        },
        "high_quality_video": {
            "crf": 23,
            "preset": "medium",
            "audio_bitrate": "128k"
        },
        "low_quality_image": {
            "scale": "20:20",
            "quality": 50
        }
    },
    "pillow": {
        "default": {
            "quality": 80,
            "optimize": True
        },
        "low_quality_image": {
            "quality": 50,
            "optimize": True
        }
    }
}


def compress_video_ffmpeg(input_file, output_dir, filename_no_ext, extension, crf=28, preset='slow', audio_bitrate='64k'):
    """使用ffmpeg压缩视频"""
    output_file = os.path.join(output_dir, f"{filename_no_ext}.{extension}")
    print(f"正在压缩视频: {input_file}")
    command = [
        'ffmpeg',
        '-i', input_file,
        '-codec:v', 'libx264',
        '-crf', str(crf),
        '-preset', preset,
        '-b:a', audio_bitrate,
        output_file
    ]
    try:
        subprocess.run(command, check=True)
        print(f"压缩完成: {os.path.basename(input_file)}")
    except subprocess.CalledProcessError as e:
        print(f"视频压缩失败 {input_file}: {e}")

def compress_image_ffmpeg(input_file, output_dir, filename_no_ext, extension, scale='40:40', quality=80, compression_level=8):
    """使用ffmpeg压缩图片并生成base64编码"""
    output_file = os.path.join(output_dir, f"{filename_no_ext}.{extension}")
    print(f"正在压缩图片: {input_file}")
    command = [
        'ffmpeg',
        '-i', input_file,
        '-vf', f"scale={scale}",
        '-q:v', str(quality),
        output_file
    ]
    try:
        subprocess.run(command, check=True)
        print(f"图片压缩完成: {os.path.basename(input_file)}")

        # 生成base64编码
        with open(output_file, 'rb') as f:
            encoded_string = base64.b64encode(f.read()).decode('utf-8')
        
        base64_output_file = os.path.join(output_dir, f"{filename_no_ext}.txt")
        with open(base64_output_file, 'w') as f:
            f.write(f"data:image/{extension};base64,{encoded_string}")
        print(f"Base64编码已保存到: {base64_output_file}")

    except subprocess.CalledProcessError as e:
        print(f"图片压缩失败 {input_file}: {e}")

def compress_image_pillow(input_file, output_dir, filename_no_ext, extension, quality=80, optimize=True):
    """使用Pillow压缩图片"""
    output_file = os.path.join(output_dir, f"{filename_no_ext}.{extension}")
    print(f"正在使用Pillow压缩图片: {input_file}")
    try:
        with Image.open(input_file) as img:
            if extension in ['jpg', 'jpeg']:
                img.save(output_file, quality=quality, optimize=optimize)
            elif extension == 'png':
                # Pillow对PNG的quality参数支持有限，主要通过optimize和compression_level控制
                img.save(output_file, optimize=optimize)
            else:
                img.save(output_file)
        print(f"Pillow图片压缩完成: {os.path.basename(input_file)}")

        # 生成base64编码
        with open(output_file, 'rb') as f:
            encoded_string = base64.b64encode(f.read()).decode('utf-8')
        
        base64_output_file = os.path.join(output_dir, f"{filename_no_ext}.txt")
        with open(base64_output_file, 'w') as f:
            f.write(f"data:image/{extension};base64,{encoded_string}")
        print(f"Base64编码已保存到: {base64_output_file}")

    except Exception as e:
        print(f"Pillow图片压缩失败 {input_file}: {e}")

def compress_image_pngquant(input_file, output_dir, filename_no_ext, extension, quality='65-80', speed=3):
    """使用pngquant压缩PNG图片
    需要安装pngquant: https://pngquant.org/
    """
    output_file = os.path.join(output_dir, f"{filename_no_ext}.{extension}")
    print(f"正在使用pngquant压缩图片: {input_file}")
    command = [
        'pngquant',
        '--quality', quality,
        '--speed', str(speed),
        '--output', output_file,
        '--force',
        input_file
    ]
    try:
        subprocess.run(command, check=True)
        print(f"pngquant图片压缩完成: {os.path.basename(input_file)}")

        # 生成base64编码
        with open(output_file, 'rb') as f:
            encoded_string = base64.b64encode(f.read()).decode('utf-8')
        
        base64_output_file = os.path.join(output_dir, f"{filename_no_ext}.txt")
        with open(base64_output_file, 'w') as f:
            f.write(f"data:image/{extension};base64,{encoded_string}")
        print(f"Base64编码已保存到: {base64_output_file}")

    except FileNotFoundError:
        print("错误: pngquant未安装或不在PATH中。请访问 https://pngquant.org/ 安装。")
    except subprocess.CalledProcessError as e:
        print(f"pngquant图片压缩失败 {input_file}: {e}")

def compress_image_jpegoptim(input_file, output_dir, filename_no_ext, extension, max_quality=80):
    """使用jpegoptim压缩JPEG图片
    需要安装jpegoptim: https://github.com/tjko/jpegoptim
    """
    output_file = os.path.join(output_dir, f"{filename_no_ext}.{extension}")
    print(f"正在使用jpegoptim压缩图片: {input_file}")
    # jpegoptim会直接修改原文件，所以先复制一份
    import shutil
    shutil.copy(input_file, output_file)

    command = [
        'jpegoptim',
        '--max', str(max_quality),
        '--strip-all',
        output_file
    ]
    try:
        subprocess.run(command, check=True)
        print(f"jpegoptim图片压缩完成: {os.path.basename(input_file)}")

        # 生成base64编码
        with open(output_file, 'rb') as f:
            encoded_string = base64.b64encode(f.read()).decode('utf-8')
        
        base64_output_file = os.path.join(output_dir, f"{filename_no_ext}.txt")
        with open(base64_output_file, 'w') as f:
            f.write(f"data:image/{extension};base64,{encoded_string}")
        print(f"Base64编码已保存到: {base64_output_file}")

    except FileNotFoundError:
        print("错误: jpegoptim未安装或不在PATH中。请访问 https://github.com/tjko/jpegoptim 安装。")
    except subprocess.CalledProcessError as e:
        print(f"jpegoptim图片压缩失败 {input_file}: {e}")

# def load_config(config_path='config.json'):
#     if os.path.exists(config_path):
#         try:
#             with open(config_path, 'r') as f:
#                 full_config = json.load(f)
#                 return full_config.get('compression_configs', {})
#          except json.JSONDecodeError as e:
#              print(f"错误: config.json 文件格式不正确: {e}")
#              sys.exit(1)
#       return {}

def main():
    print("压缩启动")
    #compression_configs = load_config()
    parser = argparse.ArgumentParser(description='压缩视频和图片文件。')
    parser.add_argument('input_directory', type=str, help='包含要压缩文件的输入目录。')
    parser.add_argument('--tool', type=str, default='ffmpeg', help='选择压缩工具 (ffmpeg, pillow, pngquant, jpegoptim)。')
    parser.add_argument('--config_key', type=str, default='default', help='选择config.json中工具的配置键。')

    args = parser.parse_args()

    input_dir = args.input_directory
    compress_tool = args.tool
    tool_config_key = args.config_key

    tool_configs = compression_configs.get(compress_tool, {})
    current_tool_config = tool_configs.get(tool_config_key, {})

    output_dir = os.path.join(input_dir, "output")

    if not os.path.isdir(input_dir):
        print(f"错误: 输入目录 '{input_dir}' 不存在或不是一个目录。")
        sys.exit(1)

    if not os.path.exists(output_dir):
        os.makedirs(output_dir)

    for filename in os.listdir(input_dir):
        file_path = os.path.join(input_dir, filename)
        if os.path.isfile(file_path):
            name, extension = os.path.splitext(filename)
            extension = extension.lstrip('.').lower()

            if extension in ['mp4', 'avi', 'mov']:
                if compress_tool == 'ffmpeg':
                    crf = current_tool_config.get('crf', 28)
                    preset = current_tool_config.get('preset', 'slow')
                    audio_bitrate = current_tool_config.get('audio_bitrate', '64k')
                    compress_video_ffmpeg(file_path, output_dir, name, extension, crf, preset, audio_bitrate)
                else:
                    print(f"不支持的视频压缩工具 '{compress_tool}'")
            elif extension in ['jpg', 'jpeg', 'png']:
                if compress_tool == 'ffmpeg':
                    scale = current_tool_config.get('scale', '40:40')
                    quality = current_tool_config.get('quality', 80)
                    compress_image_ffmpeg(file_path, output_dir, name, extension, scale, quality)
                elif compress_tool == 'pillow':
                    quality = current_tool_config.get('quality', 80)
                    optimize = current_tool_config.get('optimize', True)
                    compress_image_pillow(file_path, output_dir, name, extension, quality, optimize)
                elif compress_tool == 'pngquant':
                    if extension == 'png':
                        quality = current_tool_config.get('quality', '65-80')
                        speed = current_tool_config.get('speed', 3)
                        compress_image_pngquant(file_path, output_dir, name, extension, quality, speed)
                    else:
                        print(f"pngquant只支持PNG文件，跳过 {filename}")
                elif compress_tool == 'jpegoptim':
                    if extension in ['jpg', 'jpeg']:
                        max_quality = current_tool_config.get('max_quality', 80)
                        compress_image_jpegoptim(file_path, output_dir, name, extension, max_quality)
                    else:
                        print(f"jpegoptim只支持JPEG文件，跳过 {filename}")
                else:
                    print(f"不支持的图片压缩工具 '{compress_tool}'")
            else:
                print(f"跳过不支持的文件类型: {filename}")

    print("所有文件压缩完成！")

if __name__ == "__main__":
    main()